{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8e6810cc-5982-4293-bfbd-c91ef0aca204",
   "metadata": {},
   "source": [
    "<img src=\"http://developer.download.nvidia.com/notebooks/dlsw-notebooks/tensorrt_torchtrt_efficientnet/nvidia_logo.png\" width=\"90px\">\n",
    "\n",
    "# PySpark Tensorflow Inference\n",
    "\n",
    "### Flower Recognition with Keras Resnet50\n",
    "\n",
    "In this notebook, we demonstrate distribute inference with Resnet50 on the Databricks flower photos dataset.  \n",
    "From: https://docs.databricks.com/_static/notebooks/deep-learning/keras-metadata.html"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "858e3a8d",
   "metadata": {},
   "source": [
    "Note that cuFFT/cuDNN/cuBLAS registration errors are expected (as of `tf=2.17.0`) and will not affect behavior, as noted in [this issue.](https://github.com/tensorflow/tensorflow/issues/62075)  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "cf329ac8-0763-44bc-b0f6-b634b7dc480e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-02-04 14:00:35.457924: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
      "2025-02-04 14:00:35.465639: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:485] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
      "2025-02-04 14:00:35.473515: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:8454] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
      "2025-02-04 14:00:35.475792: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1452] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
      "2025-02-04 14:00:35.482106: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 AVX_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2025-02-04 14:00:35.843263: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import shutil\n",
    "import subprocess\n",
    "import time\n",
    "import json\n",
    "import pandas as pd\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "import uuid\n",
    " \n",
    "import tensorflow as tf\n",
    "from tensorflow.keras.applications.resnet50 import ResNet50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "532d562d",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.mkdir('models') if not os.path.exists('models') else None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "75175140",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.17.0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: All log messages before absl::InitializeLog() is called are written to STDERR\n",
      "I0000 00:00:1738706436.174805 3714100 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "I0000 00:00:1738706436.197467 3714100 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "I0000 00:00:1738706436.200398 3714100 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n"
     ]
    }
   ],
   "source": [
    "print(tf.__version__)\n",
    "\n",
    "# Enable GPU memory growth\n",
    "gpus = tf.config.experimental.list_physical_devices('GPU')\n",
    "if gpus:\n",
    "    try:\n",
    "        for gpu in gpus:\n",
    "            tf.config.experimental.set_memory_growth(gpu, True)\n",
    "    except RuntimeError as e:\n",
    "        print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02fe61b8",
   "metadata": {},
   "source": [
    "## PySpark"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b474339c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col, struct, pandas_udf, PandasUDFType\n",
    "from pyspark.ml.functions import predict_batch_udf\n",
    "from pyspark.sql.types import *\n",
    "from pyspark.sql import SparkSession\n",
    "from pyspark import SparkConf\n",
    "from typing import Iterator, Tuple"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e182cacb",
   "metadata": {},
   "source": [
    "Check the cluster environment to handle any platform-specific Spark configurations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "564b1d33",
   "metadata": {},
   "outputs": [],
   "source": [
    "on_databricks = os.environ.get(\"DATABRICKS_RUNTIME_VERSION\", False)\n",
    "on_dataproc = os.environ.get(\"DATAPROC_IMAGE_VERSION\", False)\n",
    "on_standalone = not (on_databricks or on_dataproc)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "016cdd0b",
   "metadata": {},
   "source": [
    "#### Create Spark Session\n",
    "\n",
    "For local standalone clusters, we'll connect to the cluster and create the Spark Session.  \n",
    "For CSP environments, Spark will either be preconfigured (Databricks) or we'll need to create the Spark Session (Dataproc)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "44d72768",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "25/02/04 14:00:36 WARN Utils: Your hostname, cb4ae00-lcedt resolves to a loopback address: 127.0.1.1; using 10.110.47.100 instead (on interface eno1)\n",
      "25/02/04 14:00:36 WARN Utils: Set SPARK_LOCAL_IP if you need to bind to another address\n",
      "Setting default log level to \"WARN\".\n",
      "To adjust logging level use sc.setLogLevel(newLevel). For SparkR, use setLogLevel(newLevel).\n",
      "25/02/04 14:00:37 WARN NativeCodeLoader: Unable to load native-hadoop library for your platform... using builtin-java classes where applicable\n"
     ]
    }
   ],
   "source": [
    "conf = SparkConf()\n",
    "\n",
    "if 'spark' not in globals():\n",
    "    if on_standalone:\n",
    "        import socket\n",
    "        \n",
    "        conda_env = os.environ.get(\"CONDA_PREFIX\")\n",
    "        hostname = socket.gethostname()\n",
    "        conf.setMaster(f\"spark://{hostname}:7077\")\n",
    "        conf.set(\"spark.pyspark.python\", f\"{conda_env}/bin/python\")\n",
    "        conf.set(\"spark.pyspark.driver.python\", f\"{conda_env}/bin/python\")\n",
    "    elif on_dataproc:\n",
    "        conf.set(\"spark.executorEnv.TF_GPU_ALLOCATOR\", \"cuda_malloc_async\")\n",
    "\n",
    "    conf.set(\"spark.executor.cores\", \"8\")\n",
    "    conf.set(\"spark.task.resource.gpu.amount\", \"0.125\")\n",
    "    conf.set(\"spark.executor.resource.gpu.amount\", \"1\")\n",
    "    conf.set(\"spark.sql.execution.arrow.pyspark.enabled\", \"true\")\n",
    "    conf.set(\"spark.python.worker.reuse\", \"true\")\n",
    "    conf.set(\"spark.driver.memory\", \"8g\")\n",
    "    conf.set(\"spark.executor.memory\", \"8g\")\n",
    "\n",
    "conf.set(\"spark.sql.execution.arrow.maxRecordsPerBatch\", \"512\")\n",
    "conf.set(\"spark.sql.parquet.columnarReaderBatchSize\", \"1024\")\n",
    "spark = SparkSession.builder.appName(\"spark-dl-examples\").config(conf=conf).getOrCreate()\n",
    "sc = spark.sparkContext"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61c406fa",
   "metadata": {},
   "source": [
    "Define the input and output directories."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c566dc17",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.mkdir(\"spark-dl-datasets\") if not os.path.exists(\"spark-dl-datasets\") else None\n",
    "data_path = \"spark-dl-datasets/flowers_{uuid}.parquet\".format(uuid=str(uuid.uuid1()))\n",
    "local_file_path = f\"{os.getcwd()}/{data_path}\"\n",
    "output_file_path = \"predictions/predictions\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "968d08a7-66b9-444f-b362-d8df692aef1c",
   "metadata": {},
   "source": [
    "### Prepare trained model and data for inference"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da083168-137f-492c-8769-d8f1e2111756",
   "metadata": {},
   "source": [
    "Load the ResNet-50 Model and broadcast the weights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2ddc715a-cdbc-4c49-93e9-58c9d88511da",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I0000 00:00:1738706437.771948 3714100 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "I0000 00:00:1738706437.774792 3714100 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "I0000 00:00:1738706437.777387 3714100 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "I0000 00:00:1738706437.894244 3714100 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "I0000 00:00:1738706437.895287 3714100 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "I0000 00:00:1738706437.896207 3714100 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "2025-02-04 14:00:37.897142: I tensorflow/core/common_runtime/gpu/gpu_device.cc:2021] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 40337 MB memory:  -> device: 0, name: NVIDIA RTX A6000, pci bus id: 0000:01:00.0, compute capability: 8.6\n"
     ]
    }
   ],
   "source": [
    "model = ResNet50()\n",
    "bc_model_weights = sc.broadcast(model.get_weights())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77dddfa3-e8df-4e8e-8251-64457f1ebf80",
   "metadata": {},
   "source": [
    "Load the data and save the datasets to one Parquet file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c0738bec-97d4-4946-8c49-5e6d07ff1afc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Image count: 3670\n"
     ]
    }
   ],
   "source": [
    "import pathlib\n",
    "dataset_url = \"https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz\"\n",
    "data_dir = tf.keras.utils.get_file(origin=dataset_url,\n",
    "                                   fname='flower_photos',\n",
    "                                   untar=True)\n",
    "data_dir = pathlib.Path(data_dir)\n",
    "image_count = len(list(data_dir.glob('*/*.jpg')))\n",
    "print(f\"Image count: {image_count}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d54f470a-d308-4426-8ed0-33f95155bb4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "files = [os.path.join(dp, f) for dp, dn, filenames in os.walk(data_dir) for f in filenames if os.path.splitext(f)[1] == '.jpg']\n",
    "files = files[:2048]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "64f94ee0-f1ea-47f6-a77e-be8da5d1b87a",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_data = []\n",
    "for file in files:\n",
    "    img = Image.open(file)\n",
    "    img = img.resize([224, 224])\n",
    "    data = np.asarray(img, dtype=\"float32\").reshape([224*224*3])\n",
    "\n",
    "    image_data.append({\"data\": data})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "b4ae1a98",
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.DataFrame(image_data, columns=['data']).to_parquet(data_path)\n",
    "\n",
    "if on_databricks:\n",
    "    dbutils.fs.mkdirs(\"/FileStore/spark-dl-datasets\")\n",
    "    shutil.copy(local_file_path, \"/dbfs/FileStore/{}\".format(data_path))\n",
    "    data_path = \"/dbfs/FileStore/{}\".format(data_path)\n",
    "elif on_dataproc:\n",
    "    data_dir = \"/mnt/gcs/spark-dl/spark-dl-datasets\"\n",
    "    os.mkdir(data_dir) if not os.path.exists(data_dir) else None\n",
    "    shutil.copy(local_file_path, \"/mnt/gcs/spark-dl/\" + data_path)\n",
    "    data_path = \"file:///mnt/gcs/spark-dl/\" + data_path"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2414b0f-58f2-4e4a-9d09-8ea95b38d413",
   "metadata": {},
   "source": [
    "### Save Model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "670328e3-7274-4d78-b315-487750166a3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_path = 'models/resnet50_model.keras'\n",
    "model.save(model_path)\n",
    "\n",
    "# For cloud environments, copy the model to the distributed file system.\n",
    "if on_databricks:\n",
    "    dbutils.fs.mkdirs(\"/FileStore/spark-dl-models\")\n",
    "    dbfs_model_path = \"/dbfs/FileStore/spark-dl-models/resnet50_model.keras\"\n",
    "    shutil.copy(model_path, dbfs_model_path)\n",
    "    model_path = dbfs_model_path\n",
    "elif on_dataproc:\n",
    "    # GCS is mounted at /mnt/gcs by the init script\n",
    "    models_dir = \"/mnt/gcs/spark-dl/models\"\n",
    "    os.mkdir(models_dir) if not os.path.exists(models_dir) else None\n",
    "    gcs_model_path = models_dir + \"/resnet50_model.keras\"\n",
    "    shutil.copy(model_path, gcs_model_path)\n",
    "    model_path = gcs_model_path"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b827ad56-1af0-41b7-be68-94bd203a2a70",
   "metadata": {},
   "source": [
    "### Load the data into Spark DataFrames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "8ddc22d0-b88a-4906-bd47-bf247e34feeb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2048\n"
     ]
    }
   ],
   "source": [
    "df = spark.read.parquet(data_path)\n",
    "print(df.count())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "865929b0-b016-4de4-996d-7f16176cf49c",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Model inference via Pandas UDF"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1f5a747",
   "metadata": {},
   "source": [
    "Define the function to parse the input data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "a67b3128-13c1-44f1-a0c0-7cf7a836fee3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_image(image_data):\n",
    "    image = tf.image.convert_image_dtype(\n",
    "        image_data, dtype=tf.float32) * (2. / 255) - 1\n",
    "    image = tf.reshape(image, [224, 224, 3])\n",
    "    return image"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "024e4ba2",
   "metadata": {},
   "source": [
    "Define the function for model inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "7b33185f-6d1e-4ca9-9757-fdc3d736496b",
   "metadata": {},
   "outputs": [],
   "source": [
    "@pandas_udf(ArrayType(FloatType()))\n",
    "def pandas_predict_udf(iter: Iterator[Tuple[pd.Series]]) -> Iterator[pd.Series]:\n",
    "\n",
    "    # Enable GPU memory growth to avoid CUDA OOM\n",
    "    gpus = tf.config.experimental.list_physical_devices('GPU')\n",
    "    if gpus:\n",
    "        try:\n",
    "            for gpu in gpus:\n",
    "                tf.config.experimental.set_memory_growth(gpu, True)\n",
    "        except RuntimeError as e:\n",
    "            print(e)\n",
    "\n",
    "    batch_size = 64\n",
    "    model = ResNet50(weights=None)\n",
    "    model.set_weights(bc_model_weights.value)\n",
    "    for image_batch in iter:\n",
    "        images = np.vstack(image_batch)\n",
    "        dataset = tf.data.Dataset.from_tensor_slices(images)\n",
    "        dataset = dataset.map(parse_image, num_parallel_calls=8).prefetch(\n",
    "            5000).batch(batch_size)\n",
    "        preds = model.predict(dataset)\n",
    "        yield pd.Series(list(preds))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08190547",
   "metadata": {},
   "source": [
    "Run model inference and save the results to Parquet."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "ad8c05da-db38-45ef-81d0-1f862f575ced",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 49.7 ms, sys: 17.6 ms, total: 67.3 ms\n",
      "Wall time: 15.1 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "predictions_1 = df.select(pandas_predict_udf(col(\"data\")).alias(\"prediction\"))\n",
    "results = predictions_1.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "08cb2a10",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 6:============================================>              (3 + 1) / 4]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----------------------------------------------------------------------------------------------------+\n",
      "|                                                                                          prediction|\n",
      "+----------------------------------------------------------------------------------------------------+\n",
      "|[1.2964063E-4, 2.4653607E-4, 6.7508765E-5, 1.2236452E-4, 5.7346635E-5, 3.9642912E-4, 7.033199E-6,...|\n",
      "|[4.4486973E-5, 3.5260408E-4, 4.684452E-5, 8.12069E-5, 3.179397E-5, 1.9187202E-4, 7.887208E-6, 1.3...|\n",
      "|[1.059436E-4, 2.2737762E-4, 3.0225037E-5, 6.550149E-5, 2.3658315E-5, 3.7172026E-4, 3.353684E-6, 2...|\n",
      "|[2.0393689E-5, 2.2818097E-4, 7.841931E-5, 6.991323E-5, 4.704759E-5, 9.822018E-5, 5.5858673E-6, 2....|\n",
      "|[1.13108545E-4, 2.3128217E-4, 5.283139E-5, 1.0866656E-4, 4.0229144E-5, 3.7223354E-4, 5.5677583E-6...|\n",
      "|[9.1271184E-5, 2.0681013E-4, 4.5193243E-5, 7.6812066E-5, 3.2361808E-5, 3.399333E-4, 3.8415465E-6,...|\n",
      "|[1.0792112E-4, 3.7743401E-4, 7.618583E-5, 1.24259E-4, 4.7426664E-5, 3.3307416E-4, 1.0592865E-5, 9...|\n",
      "|[2.2220212E-5, 2.7357432E-4, 3.8200575E-5, 6.235621E-5, 1.7954999E-5, 1.7249273E-4, 6.021971E-6, ...|\n",
      "|[1.1044029E-4, 2.8961376E-4, 4.2384647E-5, 1.0728626E-4, 3.0468744E-5, 4.796082E-4, 6.4537376E-6,...|\n",
      "|[9.68494E-5, 2.0567125E-4, 7.450887E-5, 1.13256065E-4, 4.609738E-5, 2.8675792E-4, 5.603957E-6, 5....|\n",
      "|[7.420906E-5, 3.2883475E-4, 1.3444667E-4, 1.7758778E-4, 8.4717096E-5, 2.2534849E-4, 1.3623082E-5,...|\n",
      "|[8.755989E-5, 2.7312606E-4, 3.59614E-5, 7.7967066E-5, 2.3571063E-5, 3.6875304E-4, 3.5629025E-6, 3...|\n",
      "|[9.7425895E-5, 2.7611412E-4, 5.74094E-5, 1.1035101E-4, 3.8303257E-5, 3.4981826E-4, 6.167147E-6, 4...|\n",
      "|[6.92996E-5, 2.5326438E-4, 5.063317E-5, 1.1494952E-4, 3.0212495E-5, 2.7857954E-4, 5.0324948E-6, 5...|\n",
      "|[4.2184765E-5, 2.4904116E-4, 1.237565E-4, 1.4271903E-4, 7.3208634E-5, 1.6054673E-4, 7.938735E-6, ...|\n",
      "|[2.719573E-5, 3.8372327E-4, 1.291892E-4, 1.5711001E-4, 7.3108524E-5, 8.553368E-5, 1.2617156E-5, 1...|\n",
      "|[3.0565643E-5, 3.55542E-4, 1.5949155E-4, 2.1368133E-4, 8.043127E-5, 1.02662845E-4, 1.3859853E-5, ...|\n",
      "|[3.311506E-5, 2.8069926E-4, 1.7956384E-4, 2.0205336E-4, 1.3665091E-4, 1.0115404E-4, 3.409792E-5, ...|\n",
      "|[4.573667E-5, 2.888326E-4, 2.3792271E-4, 2.460216E-4, 1.2164583E-4, 1.3814335E-4, 1.6352218E-5, 2...|\n",
      "|[1.2279079E-4, 2.8073761E-4, 6.365874E-5, 1.0251792E-4, 4.3527238E-5, 3.914249E-4, 8.236801E-6, 6...|\n",
      "+----------------------------------------------------------------------------------------------------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "predictions_1.show(truncate=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "40799f8e-443e-40ca-919b-391f901cb3f4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "predictions_1.write.mode(\"overwrite\").parquet(output_file_path + \"_1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7a69aa9",
   "metadata": {},
   "source": [
    "## Inference using Spark DL API\n",
    "\n",
    "Distributed inference using the PySpark [predict_batch_udf](https://spark.apache.org/docs/latest/api/python/reference/api/pyspark.ml.functions.predict_batch_udf.html#pyspark.ml.functions.predict_batch_udf):\n",
    "\n",
    "- predict_batch_fn uses Tensorflow APIs to load the model and return a predict function which operates on numpy arrays \n",
    "- predict_batch_udf will convert the Spark DataFrame columns into numpy input batches for the predict function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "dda88b46-6300-4bf7-bc10-7403f4fbbf92",
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_batch_fn():\n",
    "    import tensorflow as tf\n",
    "    from tensorflow.keras.applications.resnet50 import ResNet50\n",
    "\n",
    "    # Enable GPU memory growth\n",
    "    gpus = tf.config.experimental.list_physical_devices('GPU')\n",
    "    if gpus:\n",
    "        try:\n",
    "            for gpu in gpus:\n",
    "                tf.config.experimental.set_memory_growth(gpu, True)\n",
    "        except RuntimeError as e:\n",
    "            print(e)\n",
    "\n",
    "    model = ResNet50()\n",
    "    def predict(inputs):\n",
    "        inputs = inputs * (2. / 255) - 1\n",
    "        return model.predict(inputs)\n",
    "    return predict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "cff0e851-563d-40b6-9d05-509c22b3b7f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "classify = predict_batch_udf(predict_batch_fn,\n",
    "                             input_tensor_shapes=[[224, 224, 3]],\n",
    "                             return_type=ArrayType(FloatType()),\n",
    "                             batch_size=50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "aa7c156f-e2b3-4837-9427-ccf3a5720412",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = spark.read.parquet(data_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "80bc50ad-eaf5-4fce-a354-5e17d65e2da5",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 61.7 ms, sys: 23.1 ms, total: 84.8 ms\n",
      "Wall time: 16.7 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# first pass caches model/fn\n",
    "predictions_2 = df.select(classify(struct(\"data\")).alias(\"prediction\"))\n",
    "results = predictions_2.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "41cace80-7a4b-4929-8e63-9c83f9745e02",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 141 ms, sys: 22.2 ms, total: 163 ms\n",
      "Wall time: 16 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "predictions_2 = df.select(classify(\"data\").alias(\"prediction\"))\n",
    "results = predictions_2.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "56a2ec8a-de09-4d7c-9666-1b3c76f10657",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 52.8 ms, sys: 14.3 ms, total: 67.1 ms\n",
      "Wall time: 15.5 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "predictions_2 = df.select(classify(col(\"data\")).alias(\"prediction\"))\n",
    "results = predictions_2.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "2dcf3791",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 13:===========================================>              (3 + 1) / 4]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----------------------------------------------------------------------------------------------------+\n",
      "|                                                                                          prediction|\n",
      "+----------------------------------------------------------------------------------------------------+\n",
      "|[1.293178E-4, 2.4644283E-4, 6.760039E-5, 1.2260793E-4, 5.7431564E-5, 3.9597694E-4, 7.0522524E-6, ...|\n",
      "|[4.4487308E-5, 3.5378174E-4, 4.6667028E-5, 8.102564E-5, 3.168566E-5, 1.9189132E-4, 7.903805E-6, 1...|\n",
      "|[1.0566196E-4, 2.2684377E-4, 3.00564E-5, 6.5251304E-5, 2.3520754E-5, 3.7116173E-4, 3.331476E-6, 2...|\n",
      "|[2.0337258E-5, 2.2749524E-4, 7.8351426E-5, 6.991163E-5, 4.7081656E-5, 9.8092445E-5, 5.564894E-6, ...|\n",
      "|[1.12979564E-4, 2.3172122E-4, 5.2946547E-5, 1.0876398E-4, 4.0259067E-5, 3.7143996E-4, 5.5940513E-...|\n",
      "|[9.093228E-5, 2.0639994E-4, 4.5151268E-5, 7.666316E-5, 3.2264295E-5, 3.387436E-4, 3.832487E-6, 4....|\n",
      "|[1.0783461E-4, 3.7850672E-4, 7.660902E-5, 1.2446321E-4, 4.7591406E-5, 3.3328883E-4, 1.067249E-5, ...|\n",
      "|[2.2258617E-5, 2.7345872E-4, 3.814439E-5, 6.229726E-5, 1.79387E-5, 1.7259057E-4, 6.0371217E-6, 1....|\n",
      "|[1.1067773E-4, 2.8997674E-4, 4.2570035E-5, 1.0747747E-4, 3.0524247E-5, 4.7921995E-4, 6.489833E-6,...|\n",
      "|[9.676251E-5, 2.0588847E-4, 7.467098E-5, 1.1326933E-4, 4.6123736E-5, 2.8609246E-4, 5.627118E-6, 5...|\n",
      "|[7.4104944E-5, 3.290917E-4, 1.3448784E-4, 1.7742367E-4, 8.463227E-5, 2.2462371E-4, 1.3614881E-5, ...|\n",
      "|[8.7211796E-5, 2.7337394E-4, 3.5953894E-5, 7.7924225E-5, 2.3554327E-5, 3.67775E-4, 3.5652213E-6, ...|\n",
      "|[9.7237185E-5, 2.762026E-4, 5.7450008E-5, 1.1019135E-4, 3.831896E-5, 3.4878452E-4, 6.1574788E-6, ...|\n",
      "|[6.938849E-5, 2.5376282E-4, 5.0565883E-5, 1.14880335E-4, 3.0061366E-5, 2.7866007E-4, 5.024482E-6,...|\n",
      "|[4.2096388E-5, 2.4889092E-4, 1.2363133E-4, 1.4304162E-4, 7.337785E-5, 1.6042824E-4, 7.959722E-6, ...|\n",
      "|[2.730248E-5, 3.851789E-4, 1.293143E-4, 1.5753493E-4, 7.302161E-5, 8.547956E-5, 1.26348905E-5, 1....|\n",
      "|[3.0354899E-5, 3.5562844E-4, 1.6008675E-4, 2.1440513E-4, 8.062159E-5, 1.02023136E-4, 1.3876455E-5...|\n",
      "|[3.3083066E-5, 2.8158593E-4, 1.7979987E-4, 2.0232225E-4, 1.3704685E-4, 1.0091762E-4, 3.4243407E-5...|\n",
      "|[4.5485373E-5, 2.878148E-4, 2.3707838E-4, 2.4493985E-4, 1.21028905E-4, 1.3738636E-4, 1.6280053E-5...|\n",
      "|[1.22468E-4, 2.809503E-4, 6.3342835E-5, 1.021957E-4, 4.3373006E-5, 3.905496E-4, 8.212427E-6, 6.20...|\n",
      "+----------------------------------------------------------------------------------------------------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "predictions_2.show(truncate=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "fc511eae",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "predictions_2.write.mode(\"overwrite\").parquet(output_file_path + \"_2\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "878ca7fb",
   "metadata": {},
   "source": [
    "## Using Triton Inference Server\n",
    "In this section, we demonstrate integration with the [Triton Inference Server](https://developer.nvidia.com/nvidia-triton-inference-server), an open-source, GPU-accelerated serving solution for DL.  \n",
    "We use [PyTriton](https://github.com/triton-inference-server/pytriton), a Flask-like framework that handles client/server communication with the Triton server.  \n",
    "\n",
    "The process looks like this:\n",
    "- Distribute a PyTriton task across the Spark cluster, instructing each node to launch a Triton server process.\n",
    "- Define a Triton inference function, which contains a client that binds to the local server on a given node and sends inference requests.\n",
    "- Wrap the Triton inference function in a predict_batch_udf to launch parallel inference requests using Spark.\n",
    "- Finally, distribute a shutdown signal to terminate the Triton server processes on each node.\n",
    "\n",
    "<img src=\"../images/spark-server.png\" alt=\"drawing\" width=\"700\"/>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "2605d134-ef75-4d94-9b16-2c6d85f29bef",
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdded12d",
   "metadata": {},
   "source": [
    "Import the helper class from server_utils.py:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "a2475d41",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.addPyFile(\"server_utils.py\")\n",
    "\n",
    "from server_utils import TritonServerManager"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f6701dc",
   "metadata": {},
   "source": [
    "Define the Triton Server function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "8c8c0744-0558-4dac-bbfe-8bdde4b2af2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def triton_server(ports):\n",
    "    import time\n",
    "    import signal\n",
    "    import numpy as np\n",
    "    import tensorflow as tf\n",
    "    from tensorflow.keras.applications import ResNet50\n",
    "    from pytriton.decorators import batch\n",
    "    from pytriton.model_config import DynamicBatcher, ModelConfig, Tensor\n",
    "    from pytriton.triton import Triton, TritonConfig\n",
    "    from pyspark import TaskContext\n",
    "\n",
    "    print(f\"SERVER: Initializing ResNet on worker {TaskContext.get().partitionId()}.\")\n",
    "\n",
    "    # Enable GPU memory growth\n",
    "    gpus = tf.config.experimental.list_physical_devices('GPU')\n",
    "    if gpus:\n",
    "        try:\n",
    "            for gpu in gpus:\n",
    "                tf.config.experimental.set_memory_growth(gpu, True)\n",
    "        except RuntimeError as e:\n",
    "            print(e)\n",
    "    \n",
    "    model = ResNet50()\n",
    "    normalization_layer = tf.keras.layers.Rescaling(scale=2./255, offset=-1)\n",
    "\n",
    "    @batch\n",
    "    def _infer_fn(**inputs):\n",
    "        images = inputs[\"images\"]\n",
    "        normalized_images = normalization_layer(images)\n",
    "        return {\n",
    "            \"preds\": model.predict(normalized_images),\n",
    "        }\n",
    "\n",
    "    workspace_path = f\"/tmp/triton_{time.strftime('%m_%d_%M_%S')}\"\n",
    "    triton_conf = TritonConfig(http_port=ports[0], grpc_port=ports[1], metrics_port=ports[2])\n",
    "    with Triton(config=triton_conf, workspace=workspace_path) as triton:\n",
    "        triton.bind(\n",
    "            model_name=\"ResNet50\",\n",
    "            infer_func=_infer_fn,\n",
    "            inputs=[\n",
    "                Tensor(name=\"images\", dtype=np.float32, shape=(224, 224, 3)),\n",
    "            ],\n",
    "            outputs=[\n",
    "                Tensor(name=\"preds\", dtype=np.float32, shape=(-1,)),\n",
    "            ],\n",
    "            config=ModelConfig(\n",
    "                max_batch_size=100,\n",
    "                batcher=DynamicBatcher(max_queue_delay_microseconds=5000),  # 5ms\n",
    "            ),\n",
    "            strict=True,\n",
    "        )\n",
    "\n",
    "        def _stop_triton(signum, frame):\n",
    "            # The server manager sends SIGTERM to stop the server; this function ensures graceful cleanup.\n",
    "            print(\"SERVER: Received SIGTERM. Stopping Triton server.\")\n",
    "            triton.stop()\n",
    "\n",
    "        signal.signal(signal.SIGTERM, _stop_triton)\n",
    "\n",
    "        print(\"SERVER: Serving inference\")\n",
    "        triton.serve()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d74f7037",
   "metadata": {},
   "source": [
    "#### Start Triton servers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4bf99bde",
   "metadata": {},
   "source": [
    "The `TritonServerManager` will handle the lifecycle of Triton server instances across the Spark cluster:\n",
    "- Find available ports for HTTP/gRPC/metrics\n",
    "- Deploy a server on each node via stage-level scheduling\n",
    "- Gracefully shutdown servers across nodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2309a55c",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"ResNet50\"\n",
    "server_manager = TritonServerManager(model_name=model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "205fa1e8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-02-07 11:03:44,809 - INFO - Requesting stage-level resources: (cores=5, gpu=1.0)\n",
      "2025-02-07 11:03:44,810 - INFO - Starting 1 servers.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'cb4ae00-lcedt': (2020631, [7000, 7001, 7002])}"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Returns {'hostname', (server_pid, [http_port, grpc_port, metrics_port])}\n",
    "server_manager.start_servers(triton_server)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e49ebdbe",
   "metadata": {},
   "source": [
    "#### Define client function"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55c42174",
   "metadata": {},
   "source": [
    "Get the hostname -> url mapping from the server manager:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e4ff20e",
   "metadata": {},
   "outputs": [],
   "source": [
    "host_to_http_url = server_manager.host_to_http_url  # or server_manager.host_to_grpc_url"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "481dbd42",
   "metadata": {},
   "source": [
    "Define the Triton inference function, which returns a predict function for batch inference through the server:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "a5ab49bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def triton_fn(model_name, host_to_url):\n",
    "    import socket\n",
    "    import numpy as np\n",
    "    from pytriton.client import ModelClient\n",
    "\n",
    "    url = host_to_url[socket.gethostname()]\n",
    "    print(f\"CLIENT: Connecting to {model_name} at {url}\")\n",
    "\n",
    "    def infer_batch(inputs):\n",
    "        with ModelClient(url, model_name, inference_timeout_s=240) as client:\n",
    "            result_data = client.infer_batch(inputs)\n",
    "            return result_data[\"preds\"]\n",
    "            \n",
    "    return infer_batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "9fabcaeb-5a44-42bb-8097-5dbc2d0cee3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "classify = predict_batch_udf(partial(triton_fn, model_name=model_name, host_to_url=host_to_http_url),\n",
    "                             input_tensor_shapes=[[224, 224, 3]],\n",
    "                             return_type=ArrayType(FloatType()),\n",
    "                             batch_size=50)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcd2328e",
   "metadata": {},
   "source": [
    "#### Load DataFrame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "bbfc9009",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = spark.read.parquet(data_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c07365c-0a14-49b3-9bd8-cfb35f48b089",
   "metadata": {},
   "source": [
    "#### Run inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "e595473d-1a5d-46a6-a6ba-89d2ea903de9",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 60.9 ms, sys: 21.3 ms, total: 82.2 ms\n",
      "Wall time: 18.4 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# first pass caches model/fn\n",
    "predictions_3 = df.select(classify(struct(\"data\")).alias(\"prediction\"))\n",
    "results = predictions_3.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "5f66d468-e0b1-4589-8606-b3848063a823",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 46.3 ms, sys: 16.1 ms, total: 62.4 ms\n",
      "Wall time: 12.3 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "predictions_3 = df.select(classify(\"data\").alias(\"prediction\"))\n",
    "results = predictions_3.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "632c4c3a-fa52-4c3d-b71e-7526286e353a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 57.5 ms, sys: 16.4 ms, total: 73.9 ms\n",
      "Wall time: 12.4 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "predictions_3 = df.select(classify(col(\"data\")).alias(\"prediction\"))\n",
    "results = predictions_3.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "49870e39",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 22:===========================================>              (3 + 1) / 4]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----------------------------------------------------------------------------------------------------+\n",
      "|                                                                                          prediction|\n",
      "+----------------------------------------------------------------------------------------------------+\n",
      "|[1.293178E-4, 2.4644283E-4, 6.760039E-5, 1.2260793E-4, 5.7431564E-5, 3.9597694E-4, 7.0522524E-6, ...|\n",
      "|[4.4487308E-5, 3.5378174E-4, 4.6667028E-5, 8.102564E-5, 3.168566E-5, 1.9189132E-4, 7.903805E-6, 1...|\n",
      "|[1.0566196E-4, 2.2684377E-4, 3.00564E-5, 6.5251304E-5, 2.3520754E-5, 3.7116173E-4, 3.331476E-6, 2...|\n",
      "|[2.0337258E-5, 2.2749524E-4, 7.8351426E-5, 6.991163E-5, 4.7081656E-5, 9.8092445E-5, 5.564894E-6, ...|\n",
      "|[1.12979564E-4, 2.3172122E-4, 5.2946547E-5, 1.0876398E-4, 4.0259067E-5, 3.7143996E-4, 5.5940513E-...|\n",
      "|[9.093228E-5, 2.0639994E-4, 4.5151268E-5, 7.666316E-5, 3.2264295E-5, 3.387436E-4, 3.832487E-6, 4....|\n",
      "|[1.0783461E-4, 3.7850672E-4, 7.660902E-5, 1.2446321E-4, 4.7591406E-5, 3.3328883E-4, 1.067249E-5, ...|\n",
      "|[2.2258617E-5, 2.7345872E-4, 3.814439E-5, 6.229726E-5, 1.79387E-5, 1.7259057E-4, 6.0371217E-6, 1....|\n",
      "|[1.1067773E-4, 2.8997674E-4, 4.2570035E-5, 1.0747747E-4, 3.0524247E-5, 4.7921995E-4, 6.489833E-6,...|\n",
      "|[9.676251E-5, 2.0588847E-4, 7.467098E-5, 1.1326933E-4, 4.6123736E-5, 2.8609246E-4, 5.627118E-6, 5...|\n",
      "|[7.4104944E-5, 3.290917E-4, 1.3448784E-4, 1.7742367E-4, 8.463227E-5, 2.2462371E-4, 1.3614881E-5, ...|\n",
      "|[8.7211796E-5, 2.7337394E-4, 3.5953894E-5, 7.7924225E-5, 2.3554327E-5, 3.67775E-4, 3.5652213E-6, ...|\n",
      "|[9.7237185E-5, 2.762026E-4, 5.7450008E-5, 1.1019135E-4, 3.831896E-5, 3.4878452E-4, 6.1574788E-6, ...|\n",
      "|[6.938849E-5, 2.5376282E-4, 5.0565883E-5, 1.14880335E-4, 3.0061366E-5, 2.7866007E-4, 5.024482E-6,...|\n",
      "|[4.2096388E-5, 2.4889092E-4, 1.2363133E-4, 1.4304162E-4, 7.337785E-5, 1.6042824E-4, 7.959722E-6, ...|\n",
      "|[2.730248E-5, 3.851789E-4, 1.293143E-4, 1.5753493E-4, 7.302161E-5, 8.547956E-5, 1.26348905E-5, 1....|\n",
      "|[3.0354899E-5, 3.5562844E-4, 1.6008675E-4, 2.1440513E-4, 8.062159E-5, 1.02023136E-4, 1.3876455E-5...|\n",
      "|[3.3083066E-5, 2.8158593E-4, 1.7979987E-4, 2.0232225E-4, 1.3704685E-4, 1.0091762E-4, 3.4243407E-5...|\n",
      "|[4.5485373E-5, 2.878148E-4, 2.3707838E-4, 2.4493985E-4, 1.21028905E-4, 1.3738636E-4, 1.6280053E-5...|\n",
      "|[1.22468E-4, 2.809503E-4, 6.3342835E-5, 1.021957E-4, 4.3373006E-5, 3.905496E-4, 8.212427E-6, 6.20...|\n",
      "+----------------------------------------------------------------------------------------------------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "predictions_3.show(truncate=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "86cd59f9",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "predictions_3.write.mode(\"overwrite\").parquet(output_file_path + \"_3\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4dc06b7e-f750-40b5-9208-a035db11d937",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Stop Triton Server on each executor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "bbfcaa51-3b9f-43ff-a4a8-4b46766115b8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-02-04 14:03:34,747 - INFO - Requesting stage-level resources: (cores=5, gpu=1.0)\n",
      "2025-02-04 14:03:39,935 - INFO - Sucessfully stopped 1 servers.                 \n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[True]"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "server_manager.stop_servers()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "0d88639b-d934-4eb4-ae2f-cc13b9b10456",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not on_databricks: # on databricks, spark.stop() puts the cluster in a bad state\n",
    "    spark.stop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df8cc28a-34d7-479c-be7e-9a380d39e25e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "spark-dl-tf",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
